import os
from PIL import Image, ImageDraw, ImageFont, ImageFile
import xml.etree.ElementTree as ET


ImageFile.LOAD_TRUNCATED_IMAGES = True


def parse_xml(xml_path):
    tree = ET.parse(xml_path)
    root = tree.getroot()
    size = root.find("size")
    width = int(size.find("width").text)
    height = int(size.find("height").text)
    objects = []
    for obj in root.findall("object"):
        bndbox = obj.find("bndbox")
        xmin = int(float(bndbox.find("xmin").text))
        ymin = int(float(bndbox.find("ymin").text))
        xmax = int(float(bndbox.find("xmax").text))
        ymax = int(float(bndbox.find("ymax").text))
        name = obj.find("name").text
        objects.append((xmin, ymin, xmax, ymax, name))
    return width, height, objects


def create_new_annotation(file_name, folder, path, segmented, new_objects, output_dir):
    annotation = ET.Element("annotation")
    ET.SubElement(annotation, "folder").text = folder
    ET.SubElement(annotation, "filename").text = file_name
    ET.SubElement(annotation, "path").text = path
    source = ET.SubElement(annotation, "source")
    ET.SubElement(source, "database").text = "Unknown"
    size = ET.SubElement(annotation, "size")
    ET.SubElement(size, "width").text = str(640)
    ET.SubElement(size, "height").text = str(480)
    ET.SubElement(size, "depth").text = str(3)
    ET.SubElement(annotation, "segmented").text = str(segmented)

    for obj in new_objects:
        object_elem = ET.SubElement(annotation, "object")
        ET.SubElement(object_elem, "name").text = obj["name"]
        ET.SubElement(object_elem, "pose").text = "Unspecified"
        ET.SubElement(object_elem, "truncated").text = "0"
        ET.SubElement(object_elem, "difficult").text = "0"
        bndbox = ET.SubElement(object_elem, "bndbox")
        ET.SubElement(bndbox, "xmin").text = str(obj["xmin"])
        ET.SubElement(bndbox, "ymin").text = str(obj["ymin"])
        ET.SubElement(bndbox, "xmax").text = str(obj["xmax"])
        ET.SubElement(bndbox, "ymax").text = str(obj["ymax"])

    tree = ET.ElementTree(annotation)
    tree.write(os.path.join(output_dir, f"{os.path.splitext(file_name)[0]}.xml"))


def draw_bounding_boxes(image, objects):
    draw = ImageDraw.Draw(image)
    try:
        # Load a default font; you can specify your own font file here
        font = ImageFont.load_default()
    except IOError:
        font = None  # Fallback to a simple font if the default one is not available

    for obj in objects:
        xmin, ymin, xmax, ymax, name = obj
        draw.rectangle([xmin, ymin, xmax, ymax], outline="red", width=2)
        if font:
            bbox = draw.textbbox((0, 0), name, font=font)
            text_width = bbox[2] - bbox[0]
            text_height = bbox[3] - bbox[1]
            draw.rectangle(
                [xmin, ymin - text_height, xmin + text_width, ymin], fill="red"
            )
            draw.text((xmin, ymin - text_height), name, fill="white", font=font)
        else:
            text_width = len(name) * 10
            text_height = 15
            draw.rectangle(
                [xmin, ymin - text_height, xmin + text_width, ymin], fill="red"
            )
            draw.text((xmin, ymin - text_height), name, fill="white")

    return image


def crop_images(
    input_dir, output_dir, image_size=(640, 480), min_distance_x=320, min_distance_y=240
):
    print("input_dir:", input_dir, "\n output_dir:", output_dir)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    images_output_dir = os.path.join(output_dir, "JPEGImages")
    annotations_output_dir = os.path.join(output_dir, "Annotations")
    annotated_images_output_dir = os.path.join(output_dir, "AnnotatedImages")

    if not os.path.exists(images_output_dir):
        os.makedirs(images_output_dir)
    if not os.path.exists(annotations_output_dir):
        os.makedirs(annotations_output_dir)
    if not os.path.exists(annotated_images_output_dir):
        os.makedirs(annotated_images_output_dir)

    images_dir = os.path.join(input_dir, "JPEGImages")
    annotations_dir = os.path.join(input_dir, "Annotations")

    total_cropped_images = 0

    for filename in os.listdir(images_dir):
        if filename.endswith(".jpg") or filename.endswith(".png"):
            img_path = os.path.join(images_dir, filename)
            print("img_path:", img_path)
            xml_path = os.path.join(
                annotations_dir, os.path.splitext(filename)[0] + ".xml"
            )

            if not os.path.exists(xml_path):
                print(f"No corresponding XML file found for {img_path}")
                continue

            width, height, objects = parse_xml(xml_path)
            try:
                img = Image.open(img_path)
            except OSError as e:
                print(f"无法打开图片: {img_path} - {e}")
                continue

            # Sort objects by their distance to the center of the image
            center_x, center_y = width // 2, height // 2
            objects.sort(
                key=lambda obj: ((obj[0] + obj[2]) / 2 - center_x) ** 2
                + ((obj[1] + obj[3]) / 2 - center_y) ** 2
            )

            current_image_crops = []

            for obj in objects:
                xmin, ymin, xmax, ymax, name = obj

                # Calculate possible crop positions that include the bounding box
                left_min = max(xmin - (image_size[0] - (xmax - xmin)), 0)
                upper_min = max(ymin - (image_size[1] - (ymax - ymin)), 0)
                left_max = min(xmax, width - image_size[0])
                upper_max = min(ymax, height - image_size[1])

                for left in range(left_min, left_max + 1, min_distance_x):
                    for upper in range(upper_min, upper_max + 1, min_distance_y):
                        right = left + image_size[0]
                        lower = upper + image_size[1]

                        # Ensure the crop is within the image bounds
                        if right > width or lower > height:
                            continue

                        # Collect all objects that overlap with the current crop
                        cropped_objects = []
                        for obj in objects:
                            cxmin, cymin, cxmax, cymax, cname = obj
                            center_x, center_y = (
                                (cxmax + cxmin) / 2,
                                (cymax + cymin) / 2,
                            )
                            # print("center_x, center_y", center_x, center_y)
                            # print("left,right,upper,lower", left, right, upper, lower)

                            if (
                                center_x > left
                                and center_x < right
                                and center_y > upper
                                and center_y < lower
                            ):
                                # Adjust bounding box coordinates relative to the cropped image
                                new_xmin = max(cxmin - left, 0)
                                new_ymin = max(cymin - upper, 0)
                                new_xmax = min(cxmax - left, image_size[0])
                                new_ymax = min(cymax - upper, image_size[1])
                                cropped_objects.append(
                                    (new_xmin, new_ymin, new_xmax, new_ymax, cname)
                                )
                        # print("cropped_objects", len(cropped_objects))
                        if not cropped_objects:
                            continue

                        # Create a unique identifier for the crop
                        crop_id = (int(left), int(upper))

                        # Check if the new crop's center is far enough from existing centers
                        new_center_x = left + image_size[0] // 2
                        new_center_y = upper + image_size[1] // 2
                        valid_crop = True
                        for prev_left, prev_upper in current_image_crops:
                            prev_center_x = prev_left + image_size[0] // 2
                            prev_center_y = prev_upper + image_size[1] // 2
                            if (
                                abs(new_center_x - prev_center_x) <= min_distance_x
                                and abs(new_center_y - prev_center_y) <= min_distance_y
                            ):
                                valid_crop = False
                                break
                        # print("valid_crop", valid_crop)
                        if not valid_crop:
                            continue

                        current_image_crops.append(crop_id)

                        new_filename = f"{os.path.splitext(filename)[0]}_{int(left)}_{int(upper)}.jpg"
                        cropped_img = img.crop((left, upper, right, lower))

                        # Draw bounding boxes on the cropped image
                        cropped_img_with_boxes = draw_bounding_boxes(
                            cropped_img.copy(), cropped_objects
                        )

                        cropped_img.save(os.path.join(images_output_dir, new_filename))
                        cropped_img_with_boxes.save(
                            os.path.join(annotated_images_output_dir, new_filename)
                        )

                        create_new_annotation(
                            new_filename,
                            images_output_dir,
                            os.path.join(images_output_dir, new_filename),
                            0,
                            [
                                {
                                    "xmin": new_xmin,
                                    "ymin": new_ymin,
                                    "xmax": new_xmax,
                                    "ymax": new_ymax,
                                    "name": cname,
                                }
                                for new_xmin, new_ymin, new_xmax, new_ymax, cname in cropped_objects
                            ],
                            annotations_output_dir,
                        )

                        total_cropped_images += 1

            # Clear current image crops after processing all objects in the current image
            current_image_crops.clear()

    print(f"Total cropped images generated: {total_cropped_images}")


if __name__ == "__main__":

    # rawpath = r"F:\GitProject\yolov7-pytorch\VOCdevkit\download"
    rawpath = r"F:\GitProject\yolov7-pytorch\VOCdevkit\251003voc"
    # segpath = r"F:\GitProject\yolov7-pytorch\VOCdevkit\seg"
    segpath = r"F:\GitProject\yolov7-pytorch\VOCdevkit\251003seg"
    dirlist = os.listdir(rawpath)
    for foldername in dirlist:
        inputdir = os.path.join(rawpath, foldername)
        if os.path.isdir(inputdir):
            print(inputdir)
            crop_images(
                inputdir,
                os.path.join(segpath, foldername),
                image_size=(640, 480),
                min_distance_x=320,
                min_distance_y=240,
            )
